modelopt/torch/quantization/nn/modules/quant_module.py (1)

161-164: Make context manager re-entrant-safe by restoring prior state

If quantize_weight() is nested or called when _enable_weight_quantization was already True, hard‑resetting to False on exit can flip the state incorrectly. Preserve and restore the previous value.

     def quantize_weight(self):
         """Context in which `self.weight` is quantized."""
-        self._enable_weight_quantization = True
-        try:
-            yield
-        finally:
-            self._enable_weight_quantization = False
+        prev = getattr(self, "_enable_weight_quantization", False)
+        self._enable_weight_quantization = True
+        try:
+            yield
+        finally:
+            self._enable_weight_quantization = prev

examples/llm_qat/utils.py (1)

172-175: Perplexity value can be overwritten by existing key; compute-first merge reverses precedence

With {"perplexity": ..., **metrics}, any preexisting metrics["perplexity"] overwrites the computed value. Flip the precedence (or assign in-place) and add a small guard.

-def get_metrics_with_perplexity(metrics):
-    """Add perplexity to the metrics."""
-    metrics = {"perplexity": float(torch.exp(torch.tensor(metrics["eval_loss"]))), **metrics}
-    return metrics
+def get_metrics_with_perplexity(metrics):
+    """Add perplexity to the metrics."""
+    loss = metrics.get("eval_loss", None)
+    if loss is not None:
+        # math.exp avoids unnecessary torch tensor creation on CPU
+        import math
+        metrics["perplexity"] = float(math.exp(float(loss)))
+    return metrics

tests/_test_utils/examples/run_command.py (1)

35-44: Also set a sane default MASTER_ADDR when injecting MASTER_PORT

Some launchers expect MASTER_ADDR; defaulting it to localhost avoids env-dependent failures. Keeping allocation/race risks low is fine for tests.

 def run_example_command(cmd_parts: list[str], example_path: str, setup_free_port: bool = False):
     print(f"[{example_path}] Running command: {cmd_parts}")
     env = os.environ.copy()
 
     if setup_free_port:
         free_port = get_free_port()
         env["MASTER_PORT"] = str(free_port)
+        env.setdefault("MASTER_ADDR", "127.0.0.1")
 
     subprocess.run(cmd_parts, cwd=MODELOPT_ROOT / "examples" / example_path, env=env, check=True)

examples/llm_qat/main.py (1)

266-269: Guard eval-only printing against missing eval_loss

If trainer.evaluate() doesn’t return eval_loss, get_metrics_with_perplexity will no-op after the fix; keeping the current order is fine. Consider asserting presence during tests.

Would you like a small unit test to assert the presence/shape of evaluation metrics across backends?

modelopt/torch/utils/network.py (2)

440-454: Use isinstance for unwrapping to handle subclasses

type(model) in SUPPORTED_WRAPPERS misses subclasses. Iterate with isinstance for robustness.

-    if force_unwrap:
-        try:
-            if type(model) in SUPPORTED_WRAPPERS:
-                return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+    if force_unwrap:
+        try:
+            for wrapper_type, attr in SUPPORTED_WRAPPERS.items():
+                if isinstance(model, wrapper_type):
+                    return getattr(model, attr)
         except AttributeError:
             raise ValueError(
                 f"Model of type {type(model)} could not be forcefully unwrapped! Please manually"
                 " unwrap the model before passing it in."
             )
 
-    if type(model) in SUPPORTED_WRAPPERS:
+    for wrapper_type, attr in SUPPORTED_WRAPPERS.items():
+        if isinstance(model, wrapper_type):
             if raise_error:
                 raise ValueError(msg or f"Model {model} is wrapped by {type(model)}!")
             elif warn:
                 warnings.warn(msg or f"Model {model} is wrapped by {type(model)}; unwrapping...")
-        return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+            return getattr(model, attr)
     return model

---

599-612: Also strip DataParallel’s 'module.' prefix in get_unwrapped_name

DP inserts the same prefix as DDP; include it in the check.

-    if isinstance(model, nn.parallel.DistributedDataParallel) or (
+    if isinstance(model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)) or (
         DeepSpeedEngine is not None and isinstance(model, DeepSpeedEngine)
     ):
         name = name.removeprefix("module.")

modelopt/torch/opt/plugins/peft.py (1)

84-95: Avoid KeyError and unify with utils: use set_quantizer_state_dict.

Direct indexing quantizer_state_dict[get_unwrapped_name(name)] can KeyError on naming mismatches. Prefer the helper which tolerates missing keys.

Apply this diff:

-    if os.path.isfile(_get_quantizer_state_save_path(model_id)):
-        from modelopt.torch.quantization.nn import TensorQuantizer
-
-        quantizer_state_dict = torch.load(
-            _get_quantizer_state_save_path(model_id), map_location="cpu", weights_only=False
-        )
-        for name, module in self.named_modules():
-            if isinstance(module, TensorQuantizer):
-                module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name)])
+    if os.path.isfile(_get_quantizer_state_save_path(model_id)):
+        from modelopt.torch.quantization.utils import set_quantizer_state_dict
+        quantizer_state_dict = torch.load(
+            _get_quantizer_state_save_path(model_id), map_location="cpu", weights_only=False
+        )
+        set_quantizer_state_dict(self, quantizer_state_dict)

examples/llm_qat/simple_qat_train.py (1)

124-125: Resolve quant config from mtq.config to avoid AttributeError.

Safer to fetch from mtq.config (choices originate there).

Apply this diff:

-    model = mtq.quantize(model, getattr(mtq, args.quant_cfg), calibrate)
+    cfg = getattr(mtq.config, args.quant_cfg)
+    model = mtq.quantize(model, cfg, calibrate)

tests/examples/llm_qat/test_llm_qat.py (1)

39-45: Parametrization includes deepspeed — gate if DS isn’t installed.

Consider conditionally skipping the deepspeed case when import deepspeed fails to avoid infra‑dependent failures in CI.

modelopt/torch/quantization/plugins/transformers_trainer.py (1)

134-143: Quant config resolution: prefer mtq.config with fallback.

Some configs live under mtq.config. Use that first, then fall back to root for re‑exports.

Apply this diff:

-        if quant_args is not None and getattr(quant_args, "quant_cfg", None):
-            quant_cfg = (
-                getattr(mtq, quant_args.quant_cfg)
-                if isinstance(quant_args.quant_cfg, str)
-                else quant_args.quant_cfg
-            )
+        if quant_args is not None and getattr(quant_args, "quant_cfg", None):
+            if isinstance(quant_args.quant_cfg, str):
+                quant_cfg = getattr(getattr(mtq, "config", mtq), quant_args.quant_cfg)
+            else:
+                quant_cfg = quant_args.quant_cfg

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

examples/llm_qat/launch.sh (4)

21-26: Add a portable to_lower helper (avoid Bash 4+ dependency).

The script uses ${var,,} later, which breaks on macOS’ Bash 3.2. Add a tiny helper here and use it where you need lowercase comparisons.

Add this function just below parse_value:

to_lower() { printf '%s' "${1:-}" | tr '[:upper:]' '[:lower:]'; }

---

52-55: Fix invalid-argument error message.

Currently prints only the substring after “=”, which is confusing. Print the full flag.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument '%s'\n" "$1"

---

99-108: Lowercasing requires Bash 4+; use helper and warn when overriding backend due to --compress.

${var,,} breaks on macOS Bash. Also, silently forcing DDP when --compress True can surprise users—emit an info message.

Apply both changes (uses the to_lower helper suggested above):

-if [[ "${USE_FSDP2,,}" == "true" ]]; then
+if [[ "$(to_lower "${USE_FSDP2:-}")" == "true" ]]; then
   echo "Warning: --use_fsdp2 is deprecated. Use --backend=fsdp2 instead."
   BACKEND="fsdp2"
 fi

-# if compress is true, set backend to ddp
-if [[ "${COMPRESS,,}" == "true" ]]; then
-  BACKEND="ddp"
+# if compress is true, set backend to ddp
+if [[ "$(to_lower "${COMPRESS:-}")" == "true" ]]; then
+  if [[ "$(to_lower "${BACKEND:-}")" != "ddp" ]]; then
+    echo "Info: --compress True forces --backend=ddp (overriding '$BACKEND')."
+  fi
+  BACKEND="ddp"
 fi

---

139-142: Distillation backend check: avoid ${var,,}.

Use the helper for portability.

-  if [[ "${BACKEND,,}" == "fsdp1" || "${BACKEND,,}" == "fsdp2" ]]; then
+  if [[ "$(to_lower "${BACKEND:-}")" == "fsdp1" || "$(to_lower "${BACKEND:-}")" == "fsdp2" ]]; then
     FSDP_ARGS="$FSDP_ARGS --fsdp_cpu_ram_efficient_loading False"
   fi

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/quantization/plugins/transformers_trainer.py (1)

197-201: Bug: dataset can be None → TypeError on len(dataset)

If both train_dataset and eval_dataset are None, this crashes. Provide a clear error and compute length after selection.

Apply this diff:

-        dataset = self.train_dataset if self.train_dataset is not None else self.eval_dataset
-        num_samples = min(self.quant_args.calib_size, len(dataset))  # type: ignore [union-attr]
+        dataset = self.eval_dataset if self.eval_dataset is not None else self.train_dataset
+        assert dataset is not None, "Calibration requires either eval or train dataset."
+        num_samples = min(self.quant_args.calib_size, len(dataset))  # type: ignore [arg-type]
         dataset = torch.utils.data.Subset(dataset, list(range(num_samples)))
         data_loader = self.get_eval_dataloader(dataset)

examples/llm_qat/accelerate_config/fsdp1.yaml (1)

7-7: Activation checkpointing enabled: verify perf/compat across backends and toolchain versions.

Turning this on reduces memory but increases recompute. Please:

• Confirm torch/accelerate versions in CI support fsdp_activation_checkpointing with fsdp_use_orig_params: true and FSDP v1.
• Watch test/runtime timeouts; QAT runs may slow notably.
• Sanity‑check numerics and resume-from-checkpoint with this setting on.

Optionally gate via a CLI flag in launch.sh so users can toggle per run.

modelopt/torch/quantization/plugins/transformers_trainer.py (7)

163-169: Idempotency: avoid double‑patching accelerate.prepare

If init is called multiple times or subclasses re‑invoke the patch, _original_prepare could be overwritten. Guard the patch.

Apply this diff:

-        self._patch_accelerate_for_fsdp2_fix()
+        # Patch once
+        if not hasattr(self.accelerator, "_original_prepare"):
+            self._patch_accelerate_for_fsdp2_fix()

---

170-189: Distributed save ordering: add post‑save barrier

You barrier before saving, but not after. Add a post‑save barrier to ensure all ranks see a fully written file before proceeding.

Apply this diff:

-        if self.args.should_save:
-            torch.save(modelopt_full_state, self._modelopt_state_path)
+        if self.args.should_save:
+            torch.save(modelopt_full_state, self._modelopt_state_path)
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()

---

190-194: Robust load: add map_location to avoid device mismatches

torch.load without map_location can fail on CPU‑only envs or pick wrong GPU. Load to CPU then let restore helpers move tensors.

Apply this diff:

-        modelopt_full_state = torch.load(self._modelopt_state_path, weights_only=False)
+        modelopt_full_state = torch.load(
+            self._modelopt_state_path, map_location="cpu", weights_only=False
+        )

---

202-208: Clarify intent: use self.model by design; fix misleading comment/param

Per team learning, forward pass should invoke self.model, not the unwrapped model parameter. Rename the unused parameter and update the comment to prevent future regressions.

Apply this diff:

-        def forward_loop(model):
+        def forward_loop(_unused_unwrapped_model):
             for batch in tqdm(data_loader, desc="Calibrating"):
                 batch = self._prepare_inputs(batch)
-                # Important: We should forward pass using the unwrapped model
-                # mtq.quantize will unwrap the model pass the unwrapped model to the forward_loop
-                self.model(**batch)
+                # Important: Forward with self.model to retain wrapper hooks (DDP/FSDP2/DeepSpeed).
+                # Do not use the unwrapped model parameter passed by mtq.quantize().
+                self.model(**batch)

Note: Using the retrieved learning for this file; keeping self.model is intentional.

---

203-203: Reduce multi‑rank tqdm spam during calibration

Disable progress bars on non‑main processes.

Apply this diff:

-            for batch in tqdm(data_loader, desc="Calibrating"):
+            for batch in tqdm(
+                data_loader,
+                desc="Calibrating",
+                disable=not self.accelerator.is_local_main_process,
+            ):

---

319-321: Device placement: prefer accelerator.device over .cuda()

Calling .cuda() can select the wrong device under multi‑GPU/Accelerate. Use to(self.accelerator.device).

Apply this diff:

-        self.model.cuda()
+        self.model.to(self.accelerator.device)

---

352-354: Avoid stringly‑typed state_dict_type checks

Relying on "SHARDED_STATE_DICT" substring is brittle. Prefer comparing against the enum/type provided by the FSDP plugin if available (e.g., FSDPStateDictType.SHARDED_STATE_DICT).

If importing the enum is not feasible here, at least gate on an attribute, e.g. hasattr(self.accelerator.state.fsdp_plugin, "state_dict_type") and compare to known constants from that module.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

examples/llm_qat/accelerate_config/fsdp1.yaml (1)

23-23: num_processes: gpu — OK per earlier context; verify CI/runtime Accelerate version.

Per your note, this passes through to torchrun. Please ensure the deployed Accelerate version supports this token across all backends to avoid surprises on older runners.

examples/llm_qat/accelerate_config/deepspeed.yaml (1)

17-17: num_processes: gpu — acknowledged as intentional.

Keeping as-is per earlier discussion; just ensure runners use the compatible Accelerate version.

modelopt/torch/quantization/utils.py (1)

28-29: Right call: use get_unwrapped_name(name, model) to stabilize keys across wrappers.

modelopt/torch/quantization/plugins/transformers_trainer.py (2)

202-208: Forward pass via self.model is intentional — approved.

Per your note, use self.model(**batch) instead of the unwrapped parameter to respect Trainer hooks.

---

197-201: Fix None dataset path in calibration.

If both train and eval datasets are None, len(dataset) raises. Select dataset first and assert it exists.

-        dataset = self.train_dataset if self.train_dataset is not None else self.eval_dataset
-        num_samples = min(self.quant_args.calib_size, len(dataset))  # type: ignore [union-attr]
+        dataset = self.eval_dataset if self.eval_dataset is not None else self.train_dataset
+        assert dataset is not None, "Calibration requires either eval or train dataset."
+        num_samples = min(self.quant_args.calib_size, len(dataset))

modelopt/torch/quantization/calib/histogram.py (1)

160-163: Broaden warning to all distributed wrappers (not just DDP).

Message still ends with “DDP modules,” which re-introduces wrapper specificity. Recommend generic phrasing.

-                " method is to use the same calibration dataset across all distributed data"
-                " parallel groups so that `amax` is the same for all DDP modules."
+                " method is to use the same calibration dataset across all distributed data"
+                " parallel groups so that `amax` is the same across all ranks/modules."

modelopt/torch/quantization/nn/modules/quant_module.py (1)

161-165: Good fix: ensures cleanup even on exceptions.

The try/finally wrapper guarantees _enable_weight_quantization is reset. Consider a nested-safe counter (increment/decrement) if nested contexts are possible, otherwise this is fine.

examples/llm_qat/utils.py (1)

172-175: Avoid tensor creation; guard missing/overflow for perplexity.

Use math.exp on the float, handle missing eval_loss, and prevent overflow to keep this helper robust and cheap.

-def get_metrics_with_perplexity(metrics):
-    """Add perplexity to the metrics."""
-    metrics = {"perplexity": float(torch.exp(torch.tensor(metrics["eval_loss"]))), **metrics}
-    return metrics
+from math import exp
+
+def get_metrics_with_perplexity(metrics):
+    """Add perplexity from eval_loss if present; fall back gracefully."""
+    loss = metrics.get("eval_loss")
+    if loss is None:
+        return metrics
+    try:
+        ppl = float(exp(float(loss)))
+    except OverflowError:
+        ppl = float("inf")
+    return {**metrics, "perplexity": ppl}

tests/_test_utils/examples/run_command.py (1)

35-44: Set MASTER_ADDR alongside MASTER_PORT for deterministic local runs.

Helps avoid cross-env surprises when a global address is pre-set or missing.

 def run_example_command(cmd_parts: list[str], example_path: str, setup_free_port: bool = False):
     print(f"[{example_path}] Running command: {cmd_parts}")
     env = os.environ.copy()

     if setup_free_port:
         free_port = get_free_port()
         env["MASTER_PORT"] = str(free_port)
+        env.setdefault("MASTER_ADDR", "127.0.0.1")

     subprocess.run(cmd_parts, cwd=MODELOPT_ROOT / "examples" / example_path, env=env, check=True)

modelopt/torch/utils/network.py (1)

439-447: Improve error context when force_unwrap fails.

Include the expected attribute name in the error for faster diagnosis.

-        except AttributeError:
+        except AttributeError:
             raise ValueError(
-                f"Model of type {type(model)} could not be forcefully unwrapped! Please manually"
-                " unwrap the model before passing it in."
+                f"Model of type {type(model)} could not be forcefully unwrapped "
+                f"(missing attr '{SUPPORTED_WRAPPERS.get(type(model), 'module')}'). "
+                "Please unwrap the model before passing it in."
             )

modelopt/torch/quantization/conversion.py (1)

115-123: Avoid recomputing quantizer_state(model) twice.

Store it once to reduce traversal work on large models.

-    quantizer_state_dict = metadata["quantizer_state"]
-    unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys()
-    extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys()
+    quantizer_state_dict = metadata["quantizer_state"]
+    current_keys = quantizer_state(model).keys()
+    unmatched_keys = quantizer_state_dict.keys() - current_keys
+    extra_keys = current_keys - quantizer_state_dict.keys()

examples/llm_qat/main.py (2)

189-191: Clean stray characters in comment.

There’s a trailing “åå”.

-    # Currently useful for FSDP2 to allow for setting activation_checkpointing=True in the config file.åå
+    # Currently useful for FSDP2 to allow for setting activation_checkpointing=True in the config file.

---

266-269: Guard perplexity computation when eval_loss is absent.

Prevents KeyError on custom metrics.

-        metrics = trainer.evaluate()
-        metrics = get_metrics_with_perplexity(metrics)
-        print_rank_0(f"Evaluation results: \n{metrics}")
+        metrics = trainer.evaluate()
+        if "eval_loss" in metrics:
+            metrics = get_metrics_with_perplexity(metrics)
+        print_rank_0(f"Evaluation results: \n{metrics}")

modelopt/torch/opt/plugins/peft.py (1)

85-95: Pre-validate saved vs. current quantizer keys for clearer errors.

Direct indexing can raise a bare KeyError on mismatches. Add a short key-set check (mirrors conversion.restore_quantizer_state) for better diagnostics.

     if os.path.isfile(_get_quantizer_state_save_path(model_id)):
         from modelopt.torch.quantization.nn import TensorQuantizer
 
         quantizer_state_dict = torch.load(
             _get_quantizer_state_save_path(model_id), map_location="cpu", weights_only=False
         )
+        # Validate keys before loading for clearer messaging
+        expected = {
+            get_unwrapped_name(n, self)
+            for n, m in self.named_modules()
+            if isinstance(m, TensorQuantizer)
+        }
+        got = set(quantizer_state_dict.keys())
+        missing = expected - got
+        extra = got - expected
+        if missing or extra:
+            raise ValueError(f"Quantizer state key mismatch. missing={missing}, extra={extra}")
         for name, module in self.named_modules():
             if isinstance(module, TensorQuantizer):
-                module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name, self)])
+                key = get_unwrapped_name(name, self)
+                module.load_state_dict(quantizer_state_dict[key])

modelopt/torch/opt/conversion.py (1)

590-592: Harden restore semantics: reject wrapped models consistently.

restore() now asserts unwrapped input; restore_from_modelopt_state() does not. To avoid accidental misuse via direct calls, mirror the same check there.

 def restore_from_modelopt_state(model: ModelLike, modelopt_state: dict[str, Any]) -> nn.Module:
     """Restore the model architecture from the modelopt state dictionary based on the user-provided model."""
     # initialize ModelLikeModule if needed.
-    model = model if isinstance(model, nn.Module) else ModelLikeModule(model)
+    model = model if isinstance(model, nn.Module) else ModelLikeModule(model)
+    # Keep behavior consistent with `restore()`: do not allow wrapped models here either.
+    from modelopt.torch.utils import unwrap_model
+    model = unwrap_model(model, raise_error=True)

tests/examples/llm_qat/test_llm_qat.py (2)

39-45: Param sweep over backends is great; add env‑conditional skips for optional backends.

CI often lacks DeepSpeed (and sometimes FS*2). Guard the tests to skip when the backend isn’t available instead of failing.

 @pytest.mark.parametrize("backend", [
     "fsdp1",
     "fsdp2",
     "deepspeed",
     "ddp",
 ])
 def test_llama_qat_int4w_int8a(tiny_llama_path, tmp_path, backend):
+    if backend == "deepspeed":
+        pytest.importorskip("deepspeed")
+    if backend == "fsdp2":
+        torch = pytest.importorskip("torch")
+        from packaging.version import Version
+        if Version(torch.__version__) < Version("2.3"):
+            pytest.skip("FSDP2 requires torch>=2.3")

Repeat the small guard in test_llama_qat_int4w_int8a_direct_qat.

---

76-86: Direct QAT test: consider marking slow to keep CI wall‑time sane.

Mark as slow if your CI budget is tight.

-@pytest.mark.parametrize("backend", [
+@pytest.mark.slow
+@pytest.mark.parametrize("backend", [

modelopt/torch/quantization/utils.py (1)

459-466: Consider non‑strict load for forward/backward compat.

Older/newer TensorQuantizer shapes/keys may drift; strict=False reduces fragility while _load_from_state_dict still handles specifics.

-            module.load_state_dict(quantizer_state_dict[key])
+            module.load_state_dict(quantizer_state_dict[key], strict=False)

If you expect strict matching across versions, ignore this.

examples/llm_qat/launch.sh (2)

88-93: Quote QUANT_CFG in test to avoid word‑splitting.

Minor robustness fix.

-if [ -z $QUANT_CFG ]; then
+if [ -z "${QUANT_CFG}" ]; then
   QUANT_ARGS=""
 else
   QUANT_ARGS="--quant_cfg $QUANT_CFG --calib_size $CALIB_SIZE"
 fi

---

53-55: Error message prints the value, not the invalid flag.

Use $1 instead of ${1#*=} for clarity.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

modelopt/torch/quantization/plugins/transformers_trainer.py (3)

170-186: Potentially ineffective filter for modelopt_state_dict.

state is a (mode_str, state_dict) tuple; checking "kd_loss" in state won’t inspect metadata. The filter likely keeps everything.

-        modelopt_state["modelopt_state_dict"] = [
-            state
-            for state in modelopt_state["modelopt_state_dict"]
-            if "kd_loss" not in state and "export_student" not in state
-        ]
+        filtered = []
+        for m_str, m_state in modelopt_state["modelopt_state_dict"]:
+            meta = m_state.get("metadata", {})
+            if m_str in {"distill", "kd_loss"} or meta.get("export_student", False):
+                continue
+            filtered.append((m_str, m_state))
+        modelopt_state["modelopt_state_dict"] = filtered

Adjust the mode names if your registry uses different identifiers.

---

240-247: FSDP2 eval‑only hack looks correct; tiny nit: avoid creating opt on empty params.

If the model has no parameters (rare), next(self.model.parameters()) raises. Optional safeguard:

-            dummy_optimizer = torch.optim.SGD([next(self.model.parameters())], lr=0.0)
+            first_param = next(self.model.parameters(), None)
+            dummy_optimizer = torch.optim.SGD([first_param] if first_param is not None else [torch.nn.Parameter(torch.zeros(1, device=self.accelerator.device))], lr=0.0)

---

261-293: Accelerate FSDP2 patch is clever; leave breadcrumbs for future updates.

You’re touching private attrs (_non_persistent_buffers_set); add a one‑liner warning so future Accelerate upgrades are audited.

         def _modelopt_prepare(self, *args, **kwargs):
+            # NOTE: Relies on private Accelerate/FSDP internals; re‑verify if Accelerate/FSDP2 is updated.
             if not self.is_fsdp2:
                 return self._original_prepare(*args, **kwargs)

Please confirm the targeted Accelerate version(s) in CI where this patch is validated.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: examples/llm_qat/accelerate_config/deepspeed.yaml:17-17
Timestamp: 2025-09-15T16:36:42.871Z
Learning: In Hugging Face Accelerate configuration YAML files, num_processes can accept the string "gpu" as a value, which gets passed through to torch run under the hood and functions correctly, despite standard documentation showing integer values.

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

modelopt/torch/utils/network.py (1)

599-613: DP prefix handling added (addresses prior feedback); minor robustness nit.

Nice: now strips module. for DataParallel too. Consider stripping repeated prefixes to handle nested wrapping chains.

Apply this diff:

-    if isinstance(model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)) or (
+    if isinstance(model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)) or (
         DeepSpeedEngine is not None and isinstance(model, DeepSpeedEngine)
     ):
-        name = name.removeprefix("module.")
+        while name.startswith("module."):
+            name = name[len("module."):]

modelopt/torch/utils/network.py (2)

90-93: Align is_parallel() with SUPPORTED_WRAPPERS.

Currently it misses FSDP and DeepSpeed. Use the registry to avoid drift.

Apply this diff:

 def is_parallel(model: nn.Module) -> bool:
-    """Check if a PyTorch model is parallelized."""
-    return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))
+    """Check if a PyTorch model is wrapped by a supported parallel/distributed wrapper."""
+    return any(isinstance(model, t) for t in SUPPORTED_WRAPPERS)

---

440-454: Use isinstance instead of exact type match when unwrapping.

Exact type checks can miss subclasses/proxies (e.g., FSDP variants).

Apply this diff:

-        try:
-            if type(model) in SUPPORTED_WRAPPERS:
-                return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+        try:
+            for wrapper_t, attr in SUPPORTED_WRAPPERS.items():
+                if isinstance(model, wrapper_t):
+                    return getattr(model, attr)
         except AttributeError:
             raise ValueError(
                 f"Model of type {type(model)} could not be forcefully unwrapped! Please manually"
                 " unwrap the model before passing it in."
             )

-    if type(model) in SUPPORTED_WRAPPERS:
+    for wrapper_t, attr in SUPPORTED_WRAPPERS.items():
+        if isinstance(model, wrapper_t):
             if raise_error:
                 raise ValueError(msg or f"Model {model} is wrapped by {type(model)}!")
             elif warn:
                 warnings.warn(msg or f"Model {model} is wrapped by {type(model)}; unwrapping...")
-            return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+            return getattr(model, attr)
     return model

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/opt/dynamic.py (1)

1275-1277: Fix: pass nn.Module context to get_unwrapped_name — correct.

Using self.model ensures wrapper prefixes are stripped (DDP/DeepSpeed). Matches prior guidance.

modelopt/torch/quantization/utils.py (1)

446-456: Wrapper-stable quantizer keys — LGTM.

Using get_unwrapped_name(name, model) fixes cross-wrapper key drift. Matches prior guidance.

modelopt/torch/quantization/plugins/transformers_trainer.py (1)

197-201: Guard against missing datasets before len(...).

Order can raise TypeError when both datasets are None. Recommend:

-        dataset = self.train_dataset if self.train_dataset is not None else self.eval_dataset
-        num_samples = min(self.quant_args.calib_size, len(dataset))  # type: ignore [union-attr]
+        dataset = self.eval_dataset if self.eval_dataset is not None else self.train_dataset
+        assert dataset is not None, "Calibration requires either eval or train dataset."
+        num_samples = min(self.quant_args.calib_size, len(dataset))

modelopt/torch/quantization/nn/modules/quant_module.py (2)

157-165: Make quantize_weight re-entrant to avoid disabling on nested contexts

With the current bool flag, nested with self.quantize_weight() blocks will disable weight quantization when the inner context exits. Use a depth counter so the flag only flips to False at depth 0.

Apply:

 @contextlib.contextmanager
 def quantize_weight(self):
     """Context in which `self.weight` is quantized."""
-    self._enable_weight_quantization = True
-    try:
-        yield
-    finally:
-        self._enable_weight_quantization = False
+    depth = getattr(self, "_weight_quantization_depth", 0) + 1
+    self._weight_quantization_depth = depth
+    self._enable_weight_quantization = True
+    try:
+        yield
+    finally:
+        depth -= 1
+        self._weight_quantization_depth = depth
+        if depth == 0:
+            self._enable_weight_quantization = False

---

181-187: Register a depth counter attribute for re-entrancy

Initialize the counter alongside the existing flag.

Apply:

     def _setup(self):
         super()._setup()
         self._register_temp_attribute(
             "weight_quantizer", TensorQuantizer(self.default_quant_desc_weight)
         )
         self._register_temp_attribute("_enable_weight_quantization", False)
+        self._register_temp_attribute("_weight_quantization_depth", 0)
         self._register_dynamic_attribute("weight", self._get_quantized_weight)

examples/llm_qat/utils.py (1)

172-175: Guard against missing/invalid eval_loss and avoid unnecessary tensor ops

Handle absent eval_loss/loss keys and non-finite values; keep it torch-free for speed.

Apply:

-def get_metrics_with_perplexity(metrics):
-    """Add perplexity to the metrics."""
-    metrics = {"perplexity": float(torch.exp(torch.tensor(metrics["eval_loss"]))), **metrics}
-    return metrics
+def get_metrics_with_perplexity(metrics):
+    """Add perplexity to the metrics if loss is present."""
+    loss = metrics.get("eval_loss", metrics.get("loss"))
+    if loss is None:
+        return metrics
+    try:
+        ppl = float(torch.exp(torch.tensor(loss)))
+    except Exception:
+        return metrics
+    return {"perplexity": ppl, **metrics}

modelopt/torch/quantization/calib/histogram.py (1)

160-163: Neutral wording: not all setups are “DDP”

Since this path can run under FSDP/DeepSpeed too, avoid “DDP modules” phrasing.

Apply:

-                " method is to use the same calibration dataset across all distributed data"
-                " parallel groups so that `amax` is the same for all DDP modules."
+                " method is to use the same calibration dataset across all distributed data"
+                " parallel groups so that `amax` is consistent across all modules."

tests/_test_utils/examples/run_command.py (1)

35-44: Set MASTER_ADDR for robustness on multi-NIC hosts

Helps avoid rendezvous surprises on machines with multiple interfaces.

Apply:

 def run_example_command(cmd_parts: list[str], example_path: str, setup_free_port: bool = False):
     print(f"[{example_path}] Running command: {cmd_parts}")
     env = os.environ.copy()
 
     if setup_free_port:
         free_port = get_free_port()
         env["MASTER_PORT"] = str(free_port)
+        env.setdefault("MASTER_ADDR", "127.0.0.1")
 
     subprocess.run(cmd_parts, cwd=MODELOPT_ROOT / "examples" / example_path, env=env, check=True)

modelopt/torch/utils/network.py (3)

79-83: Narrow bare except to ImportError

Avoid swallowing unrelated runtime errors during import.

Apply:

-try:
+try:
     from deepspeed.runtime.engine import DeepSpeedEngine
-except:  # noqa: E722
+except ImportError:
     DeepSpeedEngine = None

---

438-454: Use isinstance against SUPPORTED_WRAPPERS for unwrapping (handles subclasses too)

Current type(model) in SUPPORTED_WRAPPERS misses subclasses and proxies.

Apply:

-    if force_unwrap:
-        try:
-            if type(model) in SUPPORTED_WRAPPERS:
-                return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+    if force_unwrap:
+        try:
+            for wrapper_type, attr in SUPPORTED_WRAPPERS.items():
+                if isinstance(model, wrapper_type):
+                    return getattr(model, attr)
         except AttributeError:
             raise ValueError(
                 f"Model of type {type(model)} could not be forcefully unwrapped! Please manually"
                 " unwrap the model before passing it in."
             )
 
-    if type(model) in SUPPORTED_WRAPPERS:
+    for wrapper_type, attr in SUPPORTED_WRAPPERS.items():
+        if isinstance(model, wrapper_type):
             if raise_error:
                 raise ValueError(msg or f"Model {model} is wrapped by {type(model)}!")
             elif warn:
                 warnings.warn(msg or f"Model {model} is wrapped by {type(model)}; unwrapping...")
-        return getattr(model, SUPPORTED_WRAPPERS[type(model)])
+            return getattr(model, attr)
     return model

---

90-93: is_parallel should align with SUPPORTED_WRAPPERS (include FSDP/DeepSpeed when present)

Leverage the canonical wrapper list to avoid drift.

Apply:

 def is_parallel(model: nn.Module) -> bool:
     """Check if a PyTorch model is parallelized."""
-    return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))
+    return isinstance(model, tuple(SUPPORTED_WRAPPERS.keys()))

examples/llm_qat/simple_qat_train.py (2)

118-123: Calibrate without grads and restore mode.

Avoid unnecessary autograd and keep model mode intact.

Apply this diff:

-    def calibrate(m: nn.Module):
-        for batch in calib_dataloader:
-            m(batch["input_ids"].to(device))
+    def calibrate(m: nn.Module):
+        was_training = m.training
+        m.eval()
+        with torch.no_grad():
+            for batch in calib_dataloader:
+                m(batch["input_ids"].to(device))
+        if was_training:
+            m.train()

---

109-117: Nit: avoid double .cuda()/.to(device).

Define device before model creation and move once.

Apply this diff:

-    # Load model and initialize loss
-    model = AutoModelForCausalLM.from_pretrained(args.model_path).cuda()
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
-
-    # Get dataloaders
-    train_dataloader, calib_dataloader = get_dataloader(args, tokenizer)
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Load model and initialize loss
+    model = AutoModelForCausalLM.from_pretrained(args.model_path).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+
+    # Get dataloaders
+    train_dataloader, calib_dataloader = get_dataloader(args, tokenizer)
@@
-    model.to(device)
+    # already moved above

Also applies to: 130-132

examples/llm_qat/main.py (1)

190-191: Typo: stray characters in comment.

Remove the trailing “åå”.

Apply this diff:

-    # Currently useful for FSDP2 to allow for setting activation_checkpointing=True in the config file.åå
+    # Currently useful for FSDP2 to allow for setting activation_checkpointing=True in the config file.

modelopt/torch/quantization/utils.py (1)

459-466: Make restore tolerant; early-exit on empty dict.

Loading with strict=True can fail on minor shape/version drift; also no-op fast path helps. Suggest:

-def set_quantizer_state_dict(model: nn.Module, quantizer_state_dict: dict):
+def set_quantizer_state_dict(model: nn.Module, quantizer_state_dict: dict, strict: bool = True):
     """Set the state dict of the quantizers in the model."""
     from .nn import TensorQuantizer
 
+    if not quantizer_state_dict:
+        return
     for name, module in model.named_modules():
         key = get_unwrapped_name(name, model)
         if isinstance(module, TensorQuantizer) and key in quantizer_state_dict:
-            module.load_state_dict(quantizer_state_dict[key])
+            module.load_state_dict(quantizer_state_dict[key], strict=strict)

examples/llm_qat/launch.sh (3)

52-55: Print the invalid option, not its value.

Current message shows ${1#*=} which is the value. Use $1.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

---

88-93: Quote variables in tests to avoid word-splitting.

Prevent surprises if QUANT_CFG contains spaces.

-if [ -z $QUANT_CFG ]; then
+if [ -z "${QUANT_CFG:-}" ]; then
   QUANT_ARGS=""
 else
-  QUANT_ARGS="--quant_cfg $QUANT_CFG --calib_size $CALIB_SIZE"
+  QUANT_ARGS="--quant_cfg \"$QUANT_CFG\" --calib_size \"$CALIB_SIZE\""
 fi

---

94-99: Quote MAX_STEPS check.

Avoid unary operator errors when unset.

-if [ ! -z $MAX_STEPS ]; then
+if [ -n "${MAX_STEPS:-}" ]; then
   OPTIONAL_ARGS="$OPTIONAL_ARGS --max_steps $MAX_STEPS"
 fi

modelopt/torch/quantization/plugins/transformers_trainer.py (1)

202-208: Unused param: clarify intent of forward_loop’s model arg.

Rename to underscore to avoid confusion with self.model.

-        def forward_loop(model):
+        def forward_loop(_):
             for batch in tqdm(data_loader, desc="Calibrating"):
                 batch = self._prepare_inputs(batch)
                 # Important: We should forward pass using the unwrapped model
                 # mtq.quantize will unwrap the model pass the unwrapped model to the forward_loop
                 self.model(**batch)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: examples/llm_qat/accelerate_config/deepspeed.yaml:17-17
Timestamp: 2025-09-15T16:36:42.871Z
Learning: In Hugging Face Accelerate configuration YAML files, num_processes can accept the string "gpu" as a value, which gets passed through to torch run under the hood and functions correctly, despite standard documentation showing integer values.

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

modelopt/torch/quantization/plugins/transformers_trainer.py (1)

203-210: Use of self.model in forward_loop is correct here.

Per maintainer guidance for this flow, forwarding via self.model (not the unwrapped param) is intended.

examples/llm_qat/launch.sh (3)

88-92: Quote QUANT_CFG in tests to avoid word-splitting/unset issues.

Use parameter expansion with default to prevent “[ -z ]” pitfalls when unset.

-if [ -z $QUANT_CFG ]; then
+if [ -z "${QUANT_CFG:-}" ]; then
   QUANT_ARGS=""
 else
   QUANT_ARGS="--quant_cfg $QUANT_CFG --calib_size $CALIB_SIZE"
 fi

---

95-97: Quote MAX_STEPS check.

Minor shell robustness for unset values.

-if [ ! -z $MAX_STEPS ]; then
+if [ -n "${MAX_STEPS:-}" ]; then
   OPTIONAL_ARGS="$OPTIONAL_ARGS --max_steps $MAX_STEPS"
 fi

---

52-55: Fix invalid-arg error message.

Currently prints only the value part; print the flag itself.

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

examples/llm_qat/simple_qat_train.py (1)

46-51: Optional: tune DataLoader for GPU throughput.

Pin memory and allow workers to reduce host-device stalls.

-    train_dataloader = DataLoader(
-        train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn
-    )
-    calib_dataloader = DataLoader(
-        calib_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn
-    )
+    common_dl = dict(collate_fn=collate_fn, pin_memory=torch.cuda.is_available(), num_workers=2)
+    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **common_dl)
+    calib_dataloader = DataLoader(calib_dataset, batch_size=args.batch_size, shuffle=False, **common_dl)

modelopt/torch/quantization/plugins/transformers_trainer.py (4)

241-249: Eval-only FSDP2 workaround: add safe param access.

next(self.model.parameters()) can raise StopIteration for paramless wrappers; guard defensively.

-            dummy_optimizer = torch.optim.SGD([next(self.model.parameters())], lr=0.0)
+            first_param = next((p for p in self.model.parameters()), None)
+            if first_param is None:
+                return super().evaluate(*args, **kwargs)
+            dummy_optimizer = torch.optim.SGD([first_param], lr=0.0)

---

250-261: Robust check for FULL_STATE_DICT.

State-dict type may be an enum/obj; compare via string to avoid false negatives.

-        if (
-            (not self.is_in_train)
-            and self.is_fsdp_enabled
-            and self.accelerator.state.fsdp_plugin.state_dict_type != "FULL_STATE_DICT"
-        ):
+        if (
+            (not self.is_in_train)
+            and self.is_fsdp_enabled
+            and "FULL_STATE_DICT" not in str(self.accelerator.state.fsdp_plugin.state_dict_type)
+        ):

---

262-294: Optional: guard against double-patching accelerate.prepare.

Avoid overwriting _original_prepare if already patched.

-        self.accelerator._original_prepare = self.accelerator.prepare
+        if getattr(self.accelerator, "_original_prepare", None) is None:
+            self.accelerator._original_prepare = self.accelerator.prepare
         self.accelerator.prepare = types.MethodType(_modelopt_prepare, self.accelerator)

---

321-323: Avoid unconditional .cuda() in QAD flow.

Either assert CUDA availability or move to device with fallback to improve error messaging.

-        self.model.cuda()
+        if not torch.cuda.is_available():
+            raise RuntimeError("QAD requires CUDA; no GPU detected.")
+        self.model.cuda()

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

modelopt/torch/quantization/plugins/transformers_trainer.py (2)

197-201: LGTM: dataset selection order fixed

Selecting dataset before calling len(dataset) avoids None errors previously flagged.

---

278-289: to_empty() can wipe restored quantizer buffers; snapshot/restore around prepare

Calling to_empty() before FSDP2 prepare can drop quantizer buffer storage (problematic for eval-only flows after restoring state). Snapshot quantizer state before the loop and restore it after prepare.

             tq_og_non_prsist_buffers = {}
+            saved_tq_state = get_quantizer_state_dict(model)
             for tq in (m for m in model.modules() if isinstance(m, TensorQuantizer)):
                 tq.to_empty(device=self.device)
                 tq_og_non_prsist_buffers[tq] = tq._non_persistent_buffers_set.copy()
                 tq._non_persistent_buffers_set.update(tq._buffers.keys())
@@
             outputs = self._original_prepare(*args, **kwargs)
 
             for tq in (m for m in model.modules() if isinstance(m, TensorQuantizer)):
                 tq._non_persistent_buffers_set.clear()
                 tq._non_persistent_buffers_set = tq_og_non_prsist_buffers[tq]
+            set_quantizer_state_dict(model, saved_tq_state)

Optional nit: fix the variable name typo for readability.

-            tq_og_non_prsist_buffers = {}
+            tq_og_non_persistent_buffers = {}
@@
-                tq_og_non_prsist_buffers[tq] = tq._non_persistent_buffers_set.copy()
+                tq_og_non_persistent_buffers[tq] = tq._non_persistent_buffers_set.copy()
@@
-                tq._non_persistent_buffers_set = tq_og_non_prsist_buffers[tq]
+                tq._non_persistent_buffers_set = tq_og_non_persistent_buffers[tq]

modelopt/torch/quantization/plugins/transformers_trainer.py (3)

253-260: StateDictType comparison is brittle; string vs enum mismatch

state_dict_type != "FULL_STATE_DICT" will always be True if it’s an enum. Use string containment (as done elsewhere) or compare to the enum.

-            and self.accelerator.state.fsdp_plugin.state_dict_type != "FULL_STATE_DICT"
+            and "FULL_STATE_DICT" not in str(self.accelerator.state.fsdp_plugin.state_dict_type)

---

292-293: Avoid double‑patching Accelerator.prepare

Guard _original_prepare assignment to prevent recursion if patched more than once.

-        self.accelerator._original_prepare = self.accelerator.prepare
-        self.accelerator.prepare = types.MethodType(_modelopt_prepare, self.accelerator)
+        if not hasattr(self.accelerator, "_original_prepare"):
+            self.accelerator._original_prepare = self.accelerator.prepare
+        self.accelerator.prepare = types.MethodType(_modelopt_prepare, self.accelerator)

---

203-210: Forward loop param unused; mute tqdm on non‑master and clarify comment

Rename the unused param and avoid multi‑process progress spam. Keep using self.model(**batch) as intended.

-        def forward_loop(model):
-            for batch in tqdm(data_loader, desc="Calibrating"):
+        def forward_loop(_):
+            for batch in tqdm(
+                data_loader, desc="Calibrating", disable=not self.accelerator.is_local_main_process
+            ):
                 batch = self._prepare_inputs(batch)
-                # Important: We should forward pass using the unwrapped model
-                # mtq.quantize will unwrap the model pass the unwrapped model to the forward_loop
+                # Intentionally call self.model; quantize() manages unwrapping internally.
                 self.model(**batch)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:170-189
Timestamp: 2025-09-15T20:46:28.997Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the restore_from_modelopt_state function can accept modelopt_state["modelopt_state_dict"] directly without needing to wrap it in a full dict structure or include modelopt_version.

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

examples/llm_qat/launch.sh (6)

29-52: Tighten option patterns or guard for collisions.

Patterns like --model* can match unintended flags (e.g., --model_name...). Consider exact matches to avoid accidental captures.

Example:

-    --model*)                                   MODEL=$(parse_value "$@"); [[ "$1" != *=* ]] && shift ;;
+    --model|--model=*)                           MODEL=$(parse_value "$@"); [[ "$1" != *=* ]] && shift ;;

Repeat for other options as needed.

---

88-98: Quote variables in tests to avoid word-splitting/globbing.

Unquoted vars in [ -z $QUANT_CFG ] and similar can misbehave.

Apply:

-if [ -z $QUANT_CFG ]; then
+if [ -z "${QUANT_CFG:-}" ]; then
   QUANT_ARGS=""
 else
   QUANT_ARGS="--quant_cfg $QUANT_CFG --calib_size $CALIB_SIZE"
 fi

 OPTIONAL_ARGS=""
-if [ ! -z $MAX_STEPS ]; then
+if [ -n "${MAX_STEPS:-}" ]; then
   OPTIONAL_ARGS="$OPTIONAL_ARGS --max_steps $MAX_STEPS"
 fi

---

99-107: Override to ddp on --compress: add an explicit warning.

You silently override a user-selected backend. Surface a warning to avoid surprise.

Apply:

-if [[ "${COMPRESS,,}" == "true" ]]; then
-  BACKEND="ddp"
+if [[ "${COMPRESS,,}" == "true" ]]; then
+  echo "Info: --compress enabled; forcing --backend=ddp (FSDP not supported with compression)." >&2
+  BACKEND="ddp"
 fi

---

148-179: Quote CLI argument values in CMD for safety.

Paths or model ids with spaces/shell metacharacters can break the command.

Apply:

-CMD="accelerate launch --config-file accelerate_config/$CONFIG_FILE $FSDP_ARGS \
+CMD="accelerate launch --config-file accelerate_config/$CONFIG_FILE $FSDP_ARGS \
     main.py \
-    --model_name_or_path $MODEL \
-    --model_max_length $MAX_SEQ_LENGTH \
+    --model_name_or_path \"$MODEL\" \
+    --model_max_length \"$MAX_SEQ_LENGTH\" \
     --dataloader_drop_last True \
-    --do_train $DO_TRAIN \
+    --do_train \"$DO_TRAIN\" \
     --do_eval True \
-    --output_dir $OUTPUT_DIR \
-    --dataset $DATASET \
-    --train_size $TRAIN_SIZE \
-    --eval_size $EVAL_SIZE \
-    --num_train_epochs $NUM_EPOCHS \
-    --per_device_train_batch_size $TRAIN_BS \
-    --per_device_eval_batch_size $EVAL_BS \
-    --gradient_accumulation_steps $ACCUM_STEPS \
+    --output_dir \"$OUTPUT_DIR\" \
+    --dataset \"$DATASET\" \
+    --train_size \"$TRAIN_SIZE\" \
+    --eval_size \"$EVAL_SIZE\" \
+    --num_train_epochs \"$NUM_EPOCHS\" \
+    --per_device_train_batch_size \"$TRAIN_BS\" \
+    --per_device_eval_batch_size \"$EVAL_BS\" \
+    --gradient_accumulation_steps \"$ACCUM_STEPS\" \
     --eval_accumulation_steps 1 \
     --save_strategy steps \
-    --save_steps $SAVE_STEPS \
+    --save_steps \"$SAVE_STEPS\" \
     --eval_strategy steps \
-    --eval_steps $SAVE_STEPS \
+    --eval_steps \"$SAVE_STEPS\" \
     --load_best_model_at_end True \
     --save_total_limit 2 \
-    --learning_rate $LR \
+    --learning_rate \"$LR\" \
     --weight_decay 0.0 \
     --warmup_ratio 0.1 \
     --lr_scheduler_type linear \
     --logging_steps 1 \
     --report_to tensorboard \
-    --lora $LORA \
-    --compress $COMPRESS \
+    --lora \"$LORA\" \
+    --compress \"$COMPRESS\" \
       $GRADIENT_CHECKPOINTING_ARGS $QUANT_ARGS $OPTIONAL_ARGS $DISTILLATION_ARGS
 "

---

53-55: Minor: error message should print the offending token verbatim.

${1#*=} can mangle strings; use $1.

Apply:

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

---

60-60: Trace mode: keep or gate by VERBOSE.

set -x is useful for CI but noisy for users. Consider gating on VERBOSE=1.

Example:

- set -x
+ [[ "${VERBOSE:-0}" == "1" ]] && set -x

modelopt/torch/quantization/plugins/transformers_trainer.py (4)

199-201: Add error handling for missing state file.

The method calls restore_modelopt_state_with_weights without checking if the file exists, which could raise a FileNotFoundError. This is inconsistent with the check performed in __init__ at lines 177-178.

Apply this diff to add error handling:

 def _restore_modelopt_state_with_weights(self):
+    if not os.path.exists(self._modelopt_state_path):
+        print_rank_0(f"ModelOpt state file not found: {self._modelopt_state_path}")
+        return
     restore_modelopt_state_with_weights(self.model, self._modelopt_state_path)
     print_rank_0("Restored modelopt state with weights.")

---

218-219: Consider removing the TODO for calibrate_with_adapters.

The TODO comment questions whether calibrate_with_adapters is needed, but since the PR objectives mention QLoRA support and testing, this context manager is likely required to properly disable LoRA adapters during calibration. Consider either removing the TODO or clarifying why it might not be needed in the future.

---

340-345: Add memory checks before moving model to GPU.

The self.model.cuda() call at line 339 could fail for large models that exceed GPU memory. Consider adding a try-catch or memory availability check, especially since the comment mentions that "memory efficient loading doesn't work" for QAD.

Apply this diff to add error handling:

-        self.model.cuda()
+        try:
+            self.model.cuda()
+        except torch.cuda.OutOfMemoryError as e:
+            raise RuntimeError(
+                "Failed to move model to GPU. QAD requires the entire model to fit in GPU memory. "
+                "Consider using QAT instead or reducing model/batch size."
+            ) from e

---

406-410: Add error handling for checkpoint loading failures.

The dist_cp.load_state_dict call could fail if the checkpoint is corrupted or incompatible. Consider adding error handling to provide a clearer error message.

Apply this diff to add error handling:

-    dist_cp.load_state_dict(
-        state_dict=state_dict,
-        storage_reader=dist_cp.FileSystemReader(sharded_model_path),
-        no_dist=True,
-    )
+    try:
+        dist_cp.load_state_dict(
+            state_dict=state_dict,
+            storage_reader=dist_cp.FileSystemReader(sharded_model_path),
+            no_dist=True,
+        )
+    except Exception as e:
+        print_rank_0(f"Failed to load sharded checkpoint from {sharded_model_path}: {e}")
+        raise

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:170-189
Timestamp: 2025-09-15T20:46:28.997Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the restore_from_modelopt_state function can accept modelopt_state["modelopt_state_dict"] directly without needing to wrap it in a full dict structure or include modelopt_version.

Learnt from: realAsma
PR: NVIDIA/TensorRT-Model-Optimizer#318
File: modelopt/torch/quantization/plugins/transformers_trainer.py:206-208
Timestamp: 2025-09-15T16:40:12.757Z
Learning: In modelopt/torch/quantization/plugins/transformers_trainer.py, the forward_loop function in _quantize_model should use self.model(**batch) rather than the model parameter passed to forward_loop. The model parameter should not be used for the forward pass.

examples/llm_qat/launch.sh (1)

21-25: Fix arg parsing: avoid global shift + validate missing values

Current parse_value shifts the global positional params (double-shifts with the caller) and accepts a following flag as a value. Replace with a non-shifting, validating version.

Apply:

-# Helper function to parse a single argument value
-parse_value() {
-    if [[ "$1" != *=* ]]; then shift; fi
-    echo "${1#*=}"
-}
+# Helper: extract value from "--opt=value" or "--opt value"; fail if missing (no global shift)
+parse_value() {
+  local first="$1"
+  local next="${2-}"
+  if [[ "$first" == *=* ]]; then
+    echo "${first#*=}"
+    return 0
+  fi
+  if [[ -z "$next" || "$next" == --* || "$next" == -* ]]; then
+    >&2 echo "Error: Missing value for option '$first'"
+    exit 2
+  fi
+  echo "$next"
+}

examples/llm_qat/launch.sh (6)

29-57: Correct invalid-arg error message to show the flag, not the value

Printing ${1#*=} hides the flag when using --opt=value.

Apply:

-      >&2 printf "Error: Invalid argument ${1#*=}\n"
+      >&2 printf "Error: Invalid argument %s\n" "$1"

---

88-97: Quote tests; prefer [[ -n ]] to avoid word-splitting and edge cases

Unquoted vars may misbehave when empty or starting with '-'.

Apply:

-if [ -z $QUANT_CFG ]; then
+if [[ -z "$QUANT_CFG" ]]; then
   QUANT_ARGS=""
 else
   QUANT_ARGS="--quant_cfg $QUANT_CFG --calib_size $CALIB_SIZE"
 fi

 OPTIONAL_ARGS=""
-if [ ! -z $MAX_STEPS ]; then
+if [[ -n "$MAX_STEPS" ]]; then
   OPTIONAL_ARGS="$OPTIONAL_ARGS --max_steps $MAX_STEPS"
 fi

---

99-108: Warn when COMPRESS overrides user-selected BACKEND

Currently COMPRESS silently forces ddp. Emit an explicit notice.

Apply:

 # if compress is true, set backend to ddp
 if [[ "${COMPRESS,,}" == "true" ]]; then
-  BACKEND="ddp"
+  if [[ "${BACKEND,,}" != "ddp" ]]; then
+    echo "Info: --compress enabled; overriding backend '$BACKEND' -> 'ddp'"
+  fi
+  BACKEND="ddp"
 fi

---

148-179: Quote high-risk args in CMD to handle spaces safely

Model path, output dir, and dataset may contain spaces; quote them.

Apply:

-    --model_name_or_path $MODEL \
+    --model_name_or_path "$MODEL" \
@@
-    --output_dir $OUTPUT_DIR \
+    --output_dir "$OUTPUT_DIR" \
-    --dataset $DATASET \
+    --dataset "$DATASET" \

Optional follow-up: build CMD as an array to avoid word-splitting entirely.

---

181-183: Ensure timing prints even on failure

Use a trap so elapsed time is reported when the run exits early.

Apply:

-start_time=$(date +%s)
-sh -c "$CMD"
-echo "Total time taken: $(( $(date +%s) - $start_time )) seconds"
+start_time=$(date +%s)
+trap 'echo "Total time taken: $(( $(date +%s) - start_time )) seconds"' EXIT
+sh -c "$CMD"

---

60-60: Gate shell tracing behind DEBUG

set -x can leak sensitive info; make it opt-in.

Apply:

-set -x
+[[ "${DEBUG:-0}" == "1" ]] && set -x

modelopt/torch/quantization/conversion.py (2)

115-123: Avoid recomputing quantizer_state(model) and tighten mismatch reporting.

Compute it once to reduce traversal overhead and provide clearer errors.

-    quantizer_state_dict = metadata["quantizer_state"]
-    unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys()
-    extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys()
+    quantizer_state_dict = metadata["quantizer_state"]
+    current_state = quantizer_state(model)
+    unmatched_keys = quantizer_state_dict.keys() - current_state.keys()
+    extra_keys = current_state.keys() - quantizer_state_dict.keys()

Also applies to: 124-133

---

166-172: Guard against potential unwrapped-name collisions.

If two wrapped module paths normalize to the same unwrapped key, one state can overwrite the other silently. Consider asserting uniqueness (debug-only) or logging a warning.

 def quantizer_state(model: nn.Module) -> dict[str, Any]:
     """Returns the quantizer state dict describing the quantizer states in the model."""
-    return {
+    state = {
         get_unwrapped_name(n, model): m.get_modelopt_state()
         for n, m in model.named_modules()
         if isinstance(m, (TensorQuantizer, SequentialQuantizer))
     }
+    # Optional: debug safeguard
+    # assert len(state) == len({get_unwrapped_name(n, model) for n, m in model.named_modules()
+    #                           if isinstance(m, (TensorQuantizer, SequentialQuantizer))}), \
+    #     "Duplicate unwrapped quantizer names detected."
+    return state

modelopt/torch/quantization/plugins/transformers_trainer.py (10)

163-170: Early restore/save path looks good; ensure output_dir exists.

Create the directory before save/restore to avoid surprises on fresh runs.

-        self._modelopt_state_path = os.path.join(self.args.output_dir, "modelopt_state_train.pth")
+        self._modelopt_state_path = os.path.join(self.args.output_dir, "modelopt_state_train.pth")
+        os.makedirs(self.args.output_dir, exist_ok=True)

---

190-196: Guard when modelopt_state_weights is absent.

Older checkpoints may lack weights; avoid calling set_quantizer_state_dict with None.

-        modelopt_weights = modelopt_state.pop("modelopt_state_weights", None)
+        modelopt_weights = modelopt_state.pop("modelopt_state_weights", None)
         restore_from_modelopt_state(self.model, modelopt_state)
-        set_quantizer_state_dict(self.model, modelopt_weights)
+        if modelopt_weights is not None:
+            set_quantizer_state_dict(self.model, modelopt_weights)

---

197-216: Calibration loop: limit tqdm to main process.

Prevents N progress bars under DDP/FSDP.

-        def forward_loop(model):
-            for batch in tqdm(data_loader, desc="Calibrating"):
+        def forward_loop(model):
+            for batch in tqdm(data_loader, desc="Calibrating",
+                              disable=not self.accelerator.is_main_process):
                 batch = self._prepare_inputs(batch)
                 # Important: We should forward pass using the unwrapped model
                 # mtq.quantize will unwrap the model pass the unwrapped model to the forward_loop
                 self.model(**batch)

---

212-213: Typo in TODO.

“adpaters” → “adapters”.

---

217-226: Emptying CUDA cache: gate on CUDA availability.

Avoids unnecessary call on CPU-only nodes.

-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()

---

242-249: Eval-only FSDP2 hack: OK; add small safeguard.

If model has no parameters (edge adapters), next(self.model.parameters()) raises StopIteration.

-            dummy_optimizer = torch.optim.SGD([next(self.model.parameters())], lr=0.0)
+            first_param = next(self.model.parameters(), None)
+            if first_param is None:
+                return super().evaluate(*args, **kwargs)
+            dummy_optimizer = torch.optim.SGD([first_param], lr=0.0)

---

259-274: Switching to FULL_STATE_DICT: consider restoring previous setting after save.

Prevents persistent side-effects if more saves follow.

-            self.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
-            outputs = super().save_model(*args, **kwargs)
+            prev = self.accelerator.state.fsdp_plugin.state_dict_type
+            self.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
+            try:
+                outputs = super().save_model(*args, **kwargs)
+            finally:
+                self.accelerator.state.fsdp_plugin.set_state_dict_type(prev)

---

275-307: Patch safety: avoid double-wrapping and restore non-persistent buffer set in-place.

• Don’t overwrite _original_prepare if already patched.
• Restore _non_persistent_buffers_set contents without reassigning the set (prevents reference invalidation).

-        def _modelopt_prepare(self, *args, **kwargs):
+        def _modelopt_prepare(self, *args, **kwargs):
             if not self.is_fsdp2:
                 return self._original_prepare(*args, **kwargs)
 
-            model = next((obj for obj in args if isinstance(obj, torch.nn.Module)), None)
+            model = next((obj for obj in args if isinstance(obj, torch.nn.Module)), None)
+            if model is None:
+                model = next((obj for obj in kwargs.values() if isinstance(obj, torch.nn.Module)), None)
             if model is None:
                 return self._original_prepare(*args, **kwargs)
 
             tq_og_non_prsist_buffers = {}
             for tq in (m for m in model.modules() if isinstance(m, TensorQuantizer)):
-                tq.to_empty(device=self.device)
                 tq_og_non_prsist_buffers[tq] = tq._non_persistent_buffers_set.copy()
                 tq._non_persistent_buffers_set.update(tq._buffers.keys())
 
             outputs = self._original_prepare(*args, **kwargs)
 
             for tq in (m for m in model.modules() if isinstance(m, TensorQuantizer)):
-                tq._non_persistent_buffers_set.clear()
-                tq._non_persistent_buffers_set = tq_og_non_prsist_buffers[tq]
+                tq._non_persistent_buffers_set.clear()
+                tq._non_persistent_buffers_set.update(tq_og_non_prsist_buffers[tq])
 
             return outputs
 
-        self.accelerator._original_prepare = self.accelerator.prepare
+        if getattr(self.accelerator, "_original_prepare", None) is None:
+            self.accelerator._original_prepare = self.accelerator.prepare
         self.accelerator.prepare = types.MethodType(_modelopt_prepare, self.accelerator)

---

291-296: Re-evaluate calling to_empty() on TensorQuantizer during prepare.

It can drop/restage buffers; you’re already hiding them via _non_persistent_buffers_set. If you keep it, gate it to meta-device tensors only.

---

366-375: State-dict type check by string match is brittle.

Prefer comparing the enum/constant directly if available.